package org.example.myleet.p863;

import org.example.myleet.Utils.TreeNode;

import java.util.*;

public class Solution {

    TreeNode target;

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        this.target = null;
        Map<Integer, TreeNode> parentMap = new HashMap<>();
        //先用哈希表记录所有的节点的父节点
        findParent(parentMap, new TreeNode(-1), root);
        List<Integer> result = new ArrayList<>();
        if (k == 0) {
            //distance == 0，就只有target自己
            result.add(target.val);
            return result;
        }
        //先找到target
        int i = 0;
        findTarget(root, target.val);
        Queue<TreeNode> stack1 = new LinkedList<>();
        Queue<TreeNode> stack2 = new LinkedList<>();
        Set<Integer> visited = new HashSet<>();
        //从target开始广度优先遍历各节点以寻找k距离外的节点
        stack1.offer(this.target);
        while (!stack1.isEmpty()) {
            while (!stack1.isEmpty()) {
                if (i > k) {
                    stack1.clear();
                    stack2.clear();
                    break;
                }
                TreeNode node = stack1.poll();
                if (i == k && node.val > -1) {
                    result.add(node.val);
                }
                if (node.left != null && !visited.contains(node.left.val)) {
                    stack2.offer(node.left);
                }
                if (node.right != null && !visited.contains(node.right.val)) {
                    stack2.offer(node.right);
                }
                TreeNode parent = parentMap.get(node.val);
                if (parent != null && !visited.contains(parent.val)) {
                    stack2.offer(parent);
                }
                visited.add(node.val);
            }
            stack1 = stack2;
            stack2 = new LinkedList<>();
            ++i;
        }
        return result;
    }

    private void findParent(Map<Integer, TreeNode> parentMap, TreeNode parent, TreeNode n) {
        parentMap.put(n.val, parent);
        if (n.left != null) {
            findParent(parentMap, n, n.left);
        }
        if (n.right != null) {
            findParent(parentMap, n, n.right);
        }
    }

    private void findTarget(TreeNode n, int targetValue) {
        if (n.val == targetValue) {
            this.target = n;
            return;
        }
        if (n.left != null && this.target == null) {
            findTarget(n.left, targetValue);
        }
        if (n.right != null && this.target == null) {
            findTarget(n.right, targetValue);
        }
    }
}
